#from surprise import SVD
from surprise import Dataset
from surprise import Reader
#KNNBasic：基本的协同过滤算法；KNNWithMeans：协同过滤算法的变种，考虑每个用户的平均评分；
#KNNWithZScore：协同过滤算法的变种，考虑每个用户评分的归一化操作；KNNBaseline：考虑每个用户评分的基线
from surprise import  KNNBasic,KNNWithMeans,KNNWithZScore, KNNBaseline
from surprise import accuracy
from surprise.model_selection import KFold
#import pandas as pd
from surprise import dump

# 数据读取
reader = Reader(line_format='user item rating timestamp', sep=',', skip_lines=1)
data = Dataset.load_from_file('./ratings.csv', reader=reader)

#使用相似度度量来进行评估:Pearson correlation，Cosine，Squared Distance
sim_options = {'name': 'cosine','n_epochs': 5,'user_based': True}

#k=40, min_k=1, sim_options={}, verbose=True,
#k (int整数) – 聚合时要考虑的（最大）邻居数 。 默认值为40。
#min_k (int整数) – 聚合时要考虑的最小邻居数。 如果没有足够的邻居，则将预测设置为所有评级的全局平均值。 默认值为1。
#sim_options (dict字典) – 相似性度量的选项字典。 请参阅接受选项的相似性度量配置。
#name：similarities模块中预先定义的相似度的名称。 默认值为MSD。
#user_based：选择在用户/物品间计算相似度。 对预测算法性能有巨大影响。 默认值True。
#min_support: 相似物品的最小数量 (当user-based是True) 或相似用户的最小数量 (当user-based是False)。当用户u和v评分的所有物品的数量小于最小数量，即|Iuv|<min_support时，用户u和用户v的相似度为
#shrinkage：收缩参数。（仅与Pearson correlation相似度相关）。 默认值为100

#algo = KNNBasic(sim_options=sim_options)
algo = KNNBasic()
#algo = KNNWithMeans(sim_options=sim_options)
#algo = KNNWithZScore(sim_options=sim_options)
#algo = KNNBaseline(sim_options=sim_options)

# 定义K折交叉验证迭代器，K=10
kf = KFold(n_splits=10)
for trainset, testset in kf.split(data):
    # 训练并预测
    algo.fit(trainset)
    predictions = algo.test(testset)
    # 计算RMSE
    accuracy.rmse(predictions, verbose=True)
    dump.dump('./dump_file', predictions, algo)

uid = str(196)
iid = str(302)
# 输出uid对iid的预测结果
pred = algo.predict(uid, iid, r_ui=4, verbose=True)